﻿namespace Jsl.FxCop
{
    using System;
    using System.Collections.Generic;
    using System.Diagnostics.Contracts;
    using System.Linq;
    using Microsoft.FxCop.Sdk;

    /// <summary>
    /// Builds the statement flow control for a <see cref="Method"/>.
    /// </summary>
    public class StatementFlow
    {
        private static readonly IEnumerable<Statement> emptyStatements = new Statement[] { };

        private readonly Method method;
        private Lazy<IDictionary<Statement, UniqueStatementCollection>> followingStatements;
        private Lazy<IDictionary<Statement, UniqueStatementCollection>> previousStatements;

        #region Constructors
        /// <summary>
        /// Initializes a new instance of the <see cref="StatementFlow"/> class.
        /// </summary>
        /// <param name="method">The method to get statement flow from.</param>
        public StatementFlow(Method method)
        {
            if (method == null)
            {
                throw new ArgumentNullException("method");
            }

            Contract.EndContractBlock();

            this.method = method;
            if (method.Body.Statements.Count > 0)
            {
                // Only process methods with actual statements.
                this.followingStatements = new Lazy<IDictionary<Statement, UniqueStatementCollection>>(this.BuildFollowingFlow);
                this.previousStatements = new Lazy<IDictionary<Statement, UniqueStatementCollection>>(this.BuildPreviousFlow);
            }
            else
            {
                // Skip empty methods. These are probably abstract methods since normal methods at least have a return statement in them.
                this.followingStatements = new Lazy<IDictionary<Statement, UniqueStatementCollection>>(() => new Dictionary<Statement, UniqueStatementCollection>());
                this.previousStatements = new Lazy<IDictionary<Statement, UniqueStatementCollection>>(() => new Dictionary<Statement, UniqueStatementCollection>());
            }
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="StatementFlow"/> class. Used by the <see cref="Prune"/> method.
        /// </summary>
        /// <param name="method">The method.</param>
        /// <param name="followingStatements">The following statements.</param>
        private StatementFlow(Method method, Lazy<IDictionary<Statement, UniqueStatementCollection>> followingStatements)
        {
            this.method = method;
            this.followingStatements = followingStatements;
            this.previousStatements = new Lazy<IDictionary<Statement, UniqueStatementCollection>>(this.BuildPreviousFlow);
        }
        #endregion Constructors

        #region Public Properties
        /// <summary>
        /// Gets an empty collection of statements.
        /// </summary>
        /// <value>The empty collection of statements.</value>
        public static IEnumerable<Statement> EmptyStatements
        {
            get { return StatementFlow.emptyStatements; }
        }

        /// <summary>
        /// Gets the method.
        /// </summary>
        /// <value>The method.</value>
        public Method Method
        {
            get { return this.method; }
        }

        /// <summary>
        /// Gets the all the statements in the method.
        /// </summary>
        /// <value>All the statements.</value>
        public IEnumerable<Statement> Statements
        {
            get
            {
                return this.followingStatements.Value.Keys;
            }
        }
        #endregion Public Properties

        #region Public Methods
        /// <summary>
        /// Returns all that paths that start with the specified statement.
        /// </summary>
        /// <param name="initialStatement">The initial statement.</param>
        /// <returns>All the paths that start with the specified statement.</returns>
        [System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Design", "CA1006:DoNotNestGenericTypesInMemberSignatures", Justification = "By design.")]
        public IEnumerable<IEnumerable<Statement>> AllFollowingPaths(Statement initialStatement)
        {
            return this.AllPaths(initialStatement, this.FollowingStatements);
        }

        /// <summary>
        /// Enumerates all the statements that follow this statement.
        /// </summary>
        /// <param name="statement">The statement.</param>
        /// <returns>All the statements that follow this statement.</returns>
        public IEnumerable<Statement> AllFollowingStatements(Statement statement)
        {
            return this.AllStatements(statement, this.FollowingStatements);
        }

        /// <summary>
        /// Returns all that paths that end with the specified statement.
        /// </summary>
        /// <param name="initialStatement">The initial statement.</param>
        /// <returns>All the paths that end with the specified statement.</returns>
        [System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Design", "CA1006:DoNotNestGenericTypesInMemberSignatures", Justification = "By design.")]
        public IEnumerable<IEnumerable<Statement>> AllPreviousPaths(Statement initialStatement)
        {
            return this.AllPaths(initialStatement, this.PreviousStatements);
        }

        /// <summary>
        /// Enumerates all the statements previous to this statement.
        /// </summary>
        /// <param name="statement">The statement.</param>
        /// <returns>All the statements previous to this statement.</returns>
        public IEnumerable<Statement> AllPreviousStatements(Statement statement)
        {
            return this.AllStatements(statement, this.PreviousStatements);
        }

        /// <summary>
        /// Gets the following statements.
        /// </summary>
        /// <param name="statement">The statement.</param>
        /// <returns>The following statements, if any.</returns>
        public IEnumerable<Statement> FollowingStatements(Statement statement)
        {
            Contract.Ensures(Contract.Result<IEnumerable<Statement>>() != null);

            UniqueStatementCollection statements;
            if (this.followingStatements.Value.TryGetValue(statement, out statements))
            {
                return statements;
            }
            else
            {
                return Enumerable.Empty<Statement>();
            }
        }

        /// <summary>
        /// Gets the previous statements.
        /// </summary>
        /// <param name="statement">The statement.</param>
        /// <returns>The previous statements, if any.</returns>
        public IEnumerable<Statement> PreviousStatements(Statement statement)
        {
            Contract.Ensures(Contract.Result<IEnumerable<Statement>>() != null);

            UniqueStatementCollection statements;
            if (this.previousStatements.Value.TryGetValue(statement, out statements))
            {
                return statements;
            }
            else
            {
                return Enumerable.Empty<Statement>();
            }
        }

        /// <summary>
        /// Returns a new <see cref="StatementFlow"/> with all the non-significant statements removed.
        /// </summary>
        /// <returns>A new <see cref="StatementFlow"/> with all the non-significant statements removed.</returns>
        public StatementFlow Prune()
        {
            var followingStatementBuilder = new Lazy<IDictionary<Statement, UniqueStatementCollection>>(
                    () => BuildPrunedFollowingStatements(this.followingStatements.Value));
            return new StatementFlow(this.Method, followingStatementBuilder);
        }

        /// <summary>
        /// Returns the significant following statements.
        /// </summary>
        /// <param name="statement">The statement.</param>
        /// <returns>The significant following statements.</returns>
        public IEnumerable<Statement> SignificantFollowingStatements(Statement statement)
        {
            return SignificantNextStatements(
                statement,
                nextStatement => this.followingStatements.Value[nextStatement]);
        }

        /// <summary>
        /// Returns the significant previous statements.
        /// </summary>
        /// <param name="statement">The statement.</param>
        /// <returns>The significant previous statements.</returns>
        public IEnumerable<Statement> SignificantPreviousStatements(Statement statement)
        {
            return SignificantNextStatements(
                statement,
                nextStatement => this.previousStatements.Value[nextStatement]);
        }
        #endregion Public Methods

        #region Private Methods
        private static void AddNextStatement(
            IDictionary<Statement, UniqueStatementCollection> allNextStatements,
            Statement statement,
            Statement nextStatement)
        {
            UniqueStatementCollection nextStatements = StatementFlow.EnsureStatementInCollection(allNextStatements, statement);
            nextStatements.Add(nextStatement);
        }

        private static void AddToStatements(
            HashSet<Statement> processedStatements,
            IDictionary<Statement, UniqueStatementCollection> allFollowingStatements,
            Block block)
        {
            // The first statement in the collection of statements is the FollowingStatement of this Block.
            StatementFlow.AddNextStatement(allFollowingStatements, block, block.Statements[0]);

            StatementFlow.AddToStatements(processedStatements, allFollowingStatements, block.Statements);
        }

        private static void AddToStatements(
            HashSet<Statement> processedStatements,
            IDictionary<Statement, UniqueStatementCollection> allFollowingStatements,
            Branch branch)
        {
            // The next statement is whatever the target is. If this is a "fall-through" branch then the other following
            // statement is found when processing the next higher level.
            StatementFlow.AddNextStatement(allFollowingStatements, branch, branch.Target);

            StatementFlow.AddToStatements(processedStatements, allFollowingStatements, branch.Target);
        }

        private static void AddToStatements(
            HashSet<Statement> processedStatements,
            IDictionary<Statement, UniqueStatementCollection> allFollowingStatements,
            CatchNode catchNode)
        {
            // The next statement is the catch block.
            StatementFlow.AddNextStatement(allFollowingStatements, catchNode, catchNode.Block);

            StatementFlow.AddToStatements(processedStatements, allFollowingStatements, catchNode.Block);
        }

        private static void AddToStatements(
            HashSet<Statement> processedStatements,
            IDictionary<Statement, UniqueStatementCollection> allFollowingStatements,
            FinallyNode finallyNode)
        {
            // The next statement is the finally block.
            StatementFlow.AddNextStatement(allFollowingStatements, finallyNode, finallyNode.Block);

            StatementFlow.AddToStatements(processedStatements, allFollowingStatements, finallyNode.Block);
        }

        private static void AddToStatements(
            HashSet<Statement> processedStatements,
            IDictionary<Statement, UniqueStatementCollection> allFollowingStatements,
            StatementCollection statementCollection)
        {
            for (int i = 0; i < statementCollection.Count - 1; i++)
            {
                StatementFlow.AddToStatements(processedStatements, allFollowingStatements, statementCollection[i]);

                // This statement might be something like a Block. In these cases we need to get the last statement inside
                // of this statement and hook all of it's following statements to the next statement.
                var lastStatementsOfThisStatement = StatementFlow.GetLastStatementsOf(statementCollection[i]);
                foreach (var lastStatementOfThisStatement in lastStatementsOfThisStatement)
                {
                    // Don't add anything to an unconditional branch (a goto).
                    if ((lastStatementOfThisStatement.NodeType != NodeType.Branch)
                        || (((Branch)lastStatementOfThisStatement).Condition != null))
                    {
                        StatementFlow.AddNextStatement(allFollowingStatements, lastStatementOfThisStatement, statementCollection[i + 1]);

                        StatementFlow.AddToStatements(processedStatements, allFollowingStatements, lastStatementOfThisStatement);
                    }
                }
            }

            StatementFlow.AddToStatements(processedStatements, allFollowingStatements, statementCollection.Last());
        }

        private static void AddToStatements(
            HashSet<Statement> processedStatements,
            IDictionary<Statement, UniqueStatementCollection> allFollowingStatements,
            SwitchInstruction switchStatement)
        {
            foreach (var target in switchStatement.Targets)
            {
                StatementFlow.AddNextStatement(allFollowingStatements, switchStatement, target);

                StatementFlow.AddToStatements(processedStatements, allFollowingStatements, target);
            }
        }

        private static void AddToStatements(
            HashSet<Statement> processedStatements,
            IDictionary<Statement, UniqueStatementCollection> allFollowingStatements,
            Statement statement)
        {
            if (processedStatements.Contains(statement))
            {
                return;
            }

            processedStatements.Add(statement);

            switch (statement.NodeType)
            {
                case NodeType.Block:
                    StatementFlow.AddToStatements(processedStatements, allFollowingStatements, (Block)statement);
                    break;

                case NodeType.Branch:
                    StatementFlow.AddToStatements(processedStatements, allFollowingStatements, (Branch)statement);
                    break;

                case NodeType.Catch:
                    StatementFlow.AddToStatements(processedStatements, allFollowingStatements, (CatchNode)statement);
                    break;

                case NodeType.Finally:
                    StatementFlow.AddToStatements(processedStatements, allFollowingStatements, (FinallyNode)statement);
                    break;

                case NodeType.SwitchInstruction:
                    StatementFlow.AddToStatements(processedStatements, allFollowingStatements, (SwitchInstruction)statement);
                    break;

                case NodeType.Try:
                    StatementFlow.AddToStatements(processedStatements, allFollowingStatements, (TryNode)statement);
                    break;

                default:
                    StatementFlow.EnsureStatementInCollection(allFollowingStatements, statement);
                    break;
            }
        }

        private static void AddToStatements(
            HashSet<Statement> processedStatements,
            IDictionary<Statement, UniqueStatementCollection> allFollowingStatements,
            TryNode tryNode)
        {
            StatementFlow.AddNextStatement(allFollowingStatements, tryNode, tryNode.Block);
            StatementFlow.AddToStatements(processedStatements, allFollowingStatements, tryNode.Block);

            foreach (var catchNode in tryNode.Catchers)
            {
                StatementFlow.AddNextStatement(allFollowingStatements, tryNode, catchNode);
                StatementFlow.AddToStatements(processedStatements, allFollowingStatements, catchNode);
            }

            if (tryNode.Finally != null)
            {
                StatementFlow.AddNextStatement(allFollowingStatements, tryNode, tryNode.Finally);
                StatementFlow.AddToStatements(processedStatements, allFollowingStatements, tryNode.Finally);

                foreach (var lastStatementOfBlock in GetLastStatementsOf(tryNode.Block))
                {
                    StatementFlow.AddToStatements(processedStatements, allFollowingStatements, lastStatementOfBlock);

                    // If this is a branch that is leaving the exception handler, then the branch goes through the
                    // finally first. Clear the branch target since it gets to that through the finally.
                    if ((lastStatementOfBlock.NodeType == NodeType.Branch)
                        && ((Branch)lastStatementOfBlock).LeavesExceptionBlock)
                    {
                        ClearNextStatements(allFollowingStatements, lastStatementOfBlock);
                    }

                    StatementFlow.AddNextStatement(allFollowingStatements, lastStatementOfBlock, tryNode.Finally);
                }
            }
        }

        private static IDictionary<Statement, UniqueStatementCollection> BuildPrunedFollowingStatements(
            IDictionary<Statement, UniqueStatementCollection> unprunedFollowingStatements)
        {
            var significantStatements =
                from statement in unprunedFollowingStatements.Keys
                where StatementFlow.IsSignificantStatement(statement)
                select statement;

            var allFollowingStatements = new Dictionary<Statement, UniqueStatementCollection>();
            foreach (var statement in significantStatements)
            {
                // Add this statement to the collection of statements.
                var followingStatements = StatementFlow.EnsureStatementInCollection(allFollowingStatements, statement);

                // Add all that significant statements that follow the statement.
                foreach (var followingStatement in SignificantNextStatements(statement, s => unprunedFollowingStatements[s]))
                {
                    followingStatements.Add(followingStatement);
                }
            }

            return allFollowingStatements;
        }

        private static void ClearNextStatements(
            IDictionary<Statement, UniqueStatementCollection> allNextStatements,
            Statement statement)
        {
            var statements = StatementFlow.EnsureStatementInCollection(allNextStatements, statement);
            statements.Clear();
        }

        private static UniqueStatementCollection EnsureStatementInCollection(
            IDictionary<Statement, UniqueStatementCollection> allNextStatements,
            Statement statement)
        {
            UniqueStatementCollection nextStatements;
            if (!allNextStatements.TryGetValue(statement, out nextStatements))
            {
                nextStatements = new UniqueStatementCollection();
                allNextStatements.Add(statement, nextStatements);
            }

            return nextStatements;
        }

        private static void FixTryFlow(Dictionary<Statement, UniqueStatementCollection> followingStatements)
        {
            // Notes: There is a bug in how the branches are built to a TryNode. It seems that the branch goes to
            // the Block in the TryNode instead of the TryNode itself. See the Branch_to_try_is_fixed test
            // in TestStatementFlows.cs. This might replace more than it should, but that's not a bad thing
            // for the kind of rules I'm going to write...

            var branchesToFix =
                from tryNode in followingStatements.Keys.OfType<TryNode>()
                from branch in followingStatements.Keys.OfType<Branch>()
                where followingStatements[branch].Contains(tryNode.Block)
                select new { tryNode, branch };

            foreach (var branchToFix in branchesToFix)
            {
                var statementsFollowingBranch = followingStatements[branchToFix.branch];
                statementsFollowingBranch.Remove(branchToFix.tryNode.Block);
                statementsFollowingBranch.Add(branchToFix.tryNode);
            }
        }

        private static IEnumerable<Statement> GetLastStatementsOf(Block block)
        {
            return StatementFlow.GetLastStatementsOf(block.Statements.Last());
        }

        private static IEnumerable<Statement> GetLastStatementsOf(CatchNode catchNode)
        {
            return StatementFlow.GetLastStatementsOf(catchNode.Block);
        }

        private static IEnumerable<Statement> GetLastStatementsOf(FinallyNode finallyNode)
        {
            return StatementFlow.GetLastStatementsOf(finallyNode.Block);
        }

        private static IEnumerable<Statement> GetLastStatementsOf(TryNode tryNode)
        {
            // If there is a finally node then this is always the last statement.
            if (tryNode.Finally != null)
            {
                return StatementFlow.GetLastStatementsOf(tryNode.Finally);
            }

            // If no finally nodes then the last statements are the last statements of the try block and all
            // the catch blocks.
            var endOfCatchNodes =
                from catchNode in tryNode.Catchers
                from endNode in StatementFlow.GetLastStatementsOf(catchNode)
                select endNode;
            return StatementFlow.GetLastStatementsOf(tryNode.Block)
                .Concat(endOfCatchNodes);
        }

        private static IEnumerable<Statement> GetLastStatementsOf(Statement statement)
        {
            switch (statement.NodeType)
            {
                case NodeType.Block:
                    return StatementFlow.GetLastStatementsOf((Block)statement);

                case NodeType.Catch:
                    return StatementFlow.GetLastStatementsOf((CatchNode)statement);

                case NodeType.Finally:
                    return StatementFlow.GetLastStatementsOf((FinallyNode)statement);

                case NodeType.Try:
                    return StatementFlow.GetLastStatementsOf((TryNode)statement);

                default:
                    return new Statement[] { statement };
            }
        }

        private static bool IsSignificantStatement(Statement statement)
        {
            // If this is the Pop statement that is the first statement in a catch, then ignore it.
            if (statement.NodeType == NodeType.ExpressionStatement)
            {
                if (((ExpressionStatement)statement).Expression.NodeType == NodeType.Pop)
                {
                    return false;
                }
            }

            // A branch without a condition is a goto. These can be weeded out.
            if ((statement.NodeType == NodeType.Branch)
                && (((Branch)statement).Condition == null))
            {
                return false;
            }

            return (statement.NodeType != NodeType.Nop)
                && (statement.NodeType != NodeType.Block);
        }

        private static IEnumerable<Statement> SignificantNextStatements(
            Statement statement,
            Func<Statement, IEnumerable<Statement>> nextStatements)
        {
            if (statement == null)
            {
                throw new ArgumentNullException("statement");
            }

            Contract.Ensures(Contract.Result<IEnumerable<Statement>>() != null);

            var significantStatements = new UniqueStatementCollection();
            foreach (var nextStatement in nextStatements(statement))
            {
                if (StatementFlow.IsSignificantStatement(nextStatement))
                {
                    significantStatements.Add(nextStatement);
                }
                else
                {
                    foreach (var furtherStatement in SignificantNextStatements(nextStatement, nextStatements))
                    {
                        significantStatements.Add(furtherStatement);
                    }
                }
            }

            return significantStatements;
        }

        private IEnumerable<IEnumerable<Statement>> AllPaths(
    Statement initialStatement,
    Func<Statement, IEnumerable<Statement>> getNextStatements)
        {
            // Start with an empty path and the statements after the initial statement.
            var path = new List<Statement>();
            var nextStatementsEnumerator = new Stack<IEnumerator<Statement>>();
            nextStatementsEnumerator.Push(getNextStatements(initialStatement).GetEnumerator());

            while (nextStatementsEnumerator.Count > 0)
            {
                // Get the statements we are currently working on.
                var nextStatements = nextStatementsEnumerator.Peek();

                if (nextStatements.MoveNext())
                {
                    var nextStatement = nextStatements.Current;
                    if (path.Contains(nextStatement))
                    {
                        // We looped back upon ourselves so return the current path here.
                        yield return path.ToArray();
                    }
                    else
                    {
                        // Add to the path we are currently walking.
                        path.Add(nextStatement);

                        // Get the next statements, if any.
                        var enumerable = getNextStatements(nextStatement);
                        if (enumerable.Any())
                        {
                            // There are next statements so add them to the list to process.
                            nextStatementsEnumerator.Push(enumerable.GetEnumerator());
                        }
                        else
                        {
                            // This is the end of the path so return the path here.
                            yield return path.ToArray();
                            path.RemoveAt(path.Count - 1);
                        }
                    }
                }
                else
                {
                    // We are done looking at these next statements so pop and move back a level.
                    nextStatementsEnumerator.Pop();
                    if (path.Count > 0)
                    {
                        path.RemoveAt(path.Count - 1);
                    }
                }
            }
        }

        private IEnumerable<Statement> AllStatements(Statement statement, Func<Statement, IEnumerable<Statement>> nextStatements)
        {
            var visitedStatements = new HashSet<Statement>();
            var statementsToProcess = new Stack<Statement>();
            statementsToProcess.Push(statement);

            while (statementsToProcess.Count > 0)
            {
                // Get the next statement to process. Skip it if we've already processed it.
                var currentStatement = statementsToProcess.Pop();
                foreach (var childStatement in nextStatements(currentStatement))
                {
                    // Only process statements that haven't already been seen.
                    if (!visitedStatements.Contains(childStatement))
                    {
                        // Return this statement and add it to the list that have already been processed.
                        yield return childStatement;
                        visitedStatements.Add(childStatement);

                        // Add this as another statement to follow.
                        statementsToProcess.Push(childStatement);
                    }
                }
            }
        }

        private IDictionary<Statement, UniqueStatementCollection> BuildFollowingFlow()
        {
            var followingStatements = new Dictionary<Statement, UniqueStatementCollection>();
            var processedStatements = new HashSet<Statement>();
            StatementFlow.AddToStatements(processedStatements, followingStatements, this.Method.Body.Statements);

            StatementFlow.FixTryFlow(followingStatements);

            return followingStatements;
        }

        private IDictionary<Statement, UniqueStatementCollection> BuildPreviousFlow()
        {
            var allPreviousStatements = new Dictionary<Statement, UniqueStatementCollection>();

            // Go through each statement...
            foreach (var statement in this.Statements)
            {
                StatementFlow.EnsureStatementInCollection(allPreviousStatements, statement);

                foreach (var followingStatement in this.FollowingStatements(statement))
                {
                    StatementFlow.AddNextStatement(allPreviousStatements, followingStatement, statement);
                }
            }

            return allPreviousStatements;
        }
        #endregion Private Methods
    }
}